#include "Feature.h"

#include "ExampleSet.h"

#include <iostream>
#include <fstream>		// Needed for saving and loading from files

using namespace std;

// Constants for learning tree
#define FEATURE_TESTS		100			// Number of feature tests
#define THRESHOLD_TESTS		25			// Number of threshold tests
#define MAX_TREE_DEPTH		8			// Maximum depth of tree
#define LEAF_THRESHOLD		10			// Do not split a node the number of training examples is lower than this number
#define MAX_RAM				1.7f		// Approximate free RAM in GB

// Node in a tree classifier
struct TreeNode
{
	int _tl_x;				// Top left x co-ordinate
	int _tl_y;				// Top left y co-ordinate
	int _br_x;				// Bottom right x co-ordinate
	int _br_y;				// Bottom right y co-ordinate
	featureType _f;			// Feature type
	int _o1;				// First orientation
	int _o2;				// Second orientation (use is conditional on feature type)
	float _thresh;			// Threshold for feature
	float _pos;				// Probability that a sample at this node is positive
	float _neg;				// Probability that a sample at this node is negative
	TreeNode* _l;			// Left branch of tree
	TreeNode* _r;			// Right branch of tree
};


// Class for a tree classifier
class TreeClassifier
{
	public:
		TreeClassifier();							// No argument constructor
		TreeClassifier(int sub_w, int sub_h);		// Constructor which takes sub-window dimensions
		TreeClassifier(char* directory);			// Constructs tree from text file

		// Member functions
		void SerialiseTree(char* directory);								// Saves classifier to a text function				
		void SerialiseNode(TreeNode* current, ofstream& output, int& n);	// Helper function used to recursively save the classifier to a file
		void UnserialiseNode(TreeNode*& current, ifstream& input);			// Helper function used to recursively read the classifier from a file
		void TraverseTree();												// Useful for debugging
		void TraverseNode(TreeNode* current);								// Useful for debugging
		///////////////////////////////////////////
		void LearnTree(ExampleSet* input_set);
		///////////////////////////////////////////
		void LearnTree(IntegralHistSet* input_set);							// Function to learn a decision tree
		void LearnNode(TreeNode*& current, GhostSet* g_set, int depth);		// Recursive function for learning tree
		float** RunTreeOnSameSizeSet(IntegralHistSet* imgset);				// Run the tree on a set of images that are the same size as the classifier and return the results
		TreeNode* RunTreeOnSameSizeImage(IntegralHist* hist);				// Run the tree on a single image that is the same size as the classifier and return the leaf node

		// Member variables
		int _sub_w;				// Width of the sub-window this classifier is designed to work on
		int _sub_h;				// Height of the sub-window this classifier is designed to work on
		TreeNode* _root;		// Root node of the tree
};
